"""
SimpleOrm - 基于 Tortoise ORM 的简化数据库操作类

这个模块提供了一个简化的 ORM 接口，封装了 Tortoise ORM 的常用功能，
使数据库操作更加简单和直观。

主要功能：
- 数据库连接管理
- 简化的 CRUD 操作
- 查询构建器
- 事务管理
- 批量操作
- 数据验证和序列化
"""

import asyncio
import inspect
import sys
from typing import Any, Dict, List, Optional, Type, Union, Callable
from tortoise import Tortoise, connections
from tortoise.models import Model
from tortoise.transactions import in_transaction
from tortoise.queryset import QuerySet
from tortoise.exceptions import DoesNotExist, MultipleObjectsReturned
import logging

logger = logging.getLogger(__name__)


class SimpleOrm:
    """
    基于 Tortoise ORM 的简化数据库操作类
    
    提供了简单易用的数据库操作接口，封装了 Tortoise ORM 的复杂性。
    """
    
    def __init__(self, db_url: str = None, models: List[str] = None, 
                 generate_schemas: bool = False, auto_discover: bool = True):
        """
        初始化 SimpleOrm
        
        Args:
            db_url: 数据库连接 URL
            models: 模型模块列表，如果为 None 且 auto_discover=True 则自动发现
            generate_schemas: 是否自动生成数据库表结构
            auto_discover: 是否自动发现模型类
        """
        self.db_url = db_url
        self.models = models
        self.generate_schemas = generate_schemas
        self.auto_discover = auto_discover
        self._initialized = False
        self._discovered_models = []
        
    def _discover_models(self) -> List[str]:
        """
        自动发现当前环境中的 Tortoise 模型类
        
        Returns:
            发现的模型类列表
        """
        discovered = []
        
        # 获取调用栈，找到调用 init 的模块
        frame = inspect.currentframe()
        try:
            # 向上查找调用栈，找到非 SimpleOrm 内部的调用
            while frame:
                frame = frame.f_back
                if frame and frame.f_globals.get('__name__') not in ['simpleapi.simple_orm', '__main__']:
                    caller_module = frame.f_globals.get('__name__')
                    if caller_module:
                        discovered.append(caller_module)
                        break
                elif frame and frame.f_globals.get('__name__') == '__main__':
                    # 如果是在 __main__ 中调用，扫描当前模块的所有 Model 类
                    main_models = []
                    for name, obj in frame.f_globals.items():
                        if (inspect.isclass(obj) and 
                            issubclass(obj, Model) and 
                            obj != Model and
                            not getattr(obj._meta, 'abstract', False)):
                            main_models.append(obj)
                    
                    if main_models:
                        # 将模型类直接注册到 Tortoise
                        self._discovered_models = main_models
                        return ["models"]  # 使用虚拟的 models 应用名
                    break
        finally:
            del frame
            
        # 如果没有发现，扫描所有已导入的模块
        if not discovered:
            for module_name, module in sys.modules.items():
                if module and hasattr(module, '__dict__'):
                    for attr_name in dir(module):
                        try:
                            attr = getattr(module, attr_name)
                            if (inspect.isclass(attr) and 
                                issubclass(attr, Model) and 
                                attr != Model and
                                not getattr(attr._meta, 'abstract', False)):
                                if module_name not in discovered:
                                    discovered.append(module_name)
                                break
                        except (AttributeError, TypeError):
                            continue
        
        return discovered or ["models"]
        
    async def init(self, db_url: str = None, models: List[str] = None, 
                   generate_schemas: bool = None, auto_discover: bool = None) -> None:
        """
        初始化数据库连接
        
        Args:
            db_url: 数据库连接 URL
            models: 模型模块列表，如果为 None 且 auto_discover=True 则自动发现
            generate_schemas: 是否自动生成数据库表结构
            auto_discover: 是否自动发现模型类
        """
        if self._initialized:
            logger.warning("SimpleOrm 已经初始化")
            return
            
        db_url = db_url or self.db_url
        models = models or self.models
        generate_schemas = generate_schemas if generate_schemas is not None else self.generate_schemas
        auto_discover = auto_discover if auto_discover is not None else self.auto_discover
        
        if not db_url:
            raise ValueError("必须提供数据库连接 URL")
        
        # 如果没有指定模型且启用自动发现
        if not models and auto_discover:
            models = self._discover_models()
            logger.info(f"自动发现模型: {models}")
        
        if not models:
            raise ValueError("未找到任何模型，请指定 models 参数或启用 auto_discover")
            
        config = {
            "connections": {"default": db_url},
            "apps": {
                "models": {
                    "models": models,
                    "default_connection": "default",
                }
            },
        }
        
        # 如果有直接发现的模型类，需要特殊处理
        if self._discovered_models:
            # 创建一个临时模块来存放模型
            import types
            temp_module = types.ModuleType("models")
            temp_module.__file__ = "<auto_discovered>"
            
            # 将模型类添加到临时模块
            for model_class in self._discovered_models:
                setattr(temp_module, model_class.__name__, model_class)
            
            # 将临时模块添加到 sys.modules
            sys.modules["models"] = temp_module
            
            try:
                await Tortoise.init(config)
            finally:
                # 清理临时模块（可选）
                pass
        else:
            await Tortoise.init(config)
        
        if generate_schemas:
            await Tortoise.generate_schemas()
            
        self._initialized = True
        logger.info("SimpleOrm 初始化完成")
        
    async def close(self) -> None:
        """关闭数据库连接"""
        if self._initialized:
            await Tortoise.close_connections()
            self._initialized = False
            logger.info("数据库连接已关闭")
            
    def _ensure_initialized(self) -> None:
        """确保 ORM 已初始化"""
        if not self._initialized:
            raise RuntimeError("SimpleOrm 尚未初始化，请先调用 init() 方法")
            
    # CRUD 操作
    async def create(self, model_class: Type[Model], **kwargs) -> Model:
        """
        创建新记录
        
        Args:
            model_class: 模型类
            **kwargs: 字段值
            
        Returns:
            创建的模型实例
        """
        self._ensure_initialized()
        instance = await model_class.create(**kwargs)
        logger.debug(f"创建记录: {model_class.__name__} - {instance.pk}")
        return instance
        
    async def get(self, model_class: Type[Model], **kwargs) -> Optional[Model]:
        """
        获取单个记录
        
        Args:
            model_class: 模型类
            **kwargs: 查询条件
            
        Returns:
            模型实例或 None
        """
        self._ensure_initialized()
        try:
            return await model_class.get(**kwargs)
        except DoesNotExist:
            return None
            
    async def get_or_create(self, model_class: Type[Model], defaults: Dict = None, 
                           **kwargs) -> tuple[Model, bool]:
        """
        获取或创建记录
        
        Args:
            model_class: 模型类
            defaults: 创建时的默认值
            **kwargs: 查询条件
            
        Returns:
            (模型实例, 是否为新创建)
        """
        self._ensure_initialized()
        return await model_class.get_or_create(defaults=defaults or {}, **kwargs)
        
    async def update(self, model_class: Type[Model], filter_kwargs: Dict, 
                    **update_kwargs) -> int:
        """
        更新记录
        
        Args:
            model_class: 模型类
            filter_kwargs: 过滤条件
            **update_kwargs: 更新的字段值
            
        Returns:
            更新的记录数量
        """
        self._ensure_initialized()
        count = await model_class.filter(**filter_kwargs).update(**update_kwargs)
        logger.debug(f"更新记录: {model_class.__name__} - {count} 条")
        return count
        
    async def delete(self, model_class: Type[Model], **kwargs) -> int:
        """
        删除记录
        
        Args:
            model_class: 模型类
            **kwargs: 删除条件
            
        Returns:
            删除的记录数量
        """
        self._ensure_initialized()
        count = await model_class.filter(**kwargs).delete()
        logger.debug(f"删除记录: {model_class.__name__} - {count} 条")
        return count
        
    async def list(self, model_class: Type[Model], limit: int = None, 
                  offset: int = None, order_by: List[str] = None, 
                  **kwargs) -> List[Model]:
        """
        获取记录列表
        
        Args:
            model_class: 模型类
            limit: 限制数量
            offset: 偏移量
            order_by: 排序字段
            **kwargs: 过滤条件
            
        Returns:
            模型实例列表
        """
        self._ensure_initialized()
        queryset = model_class.filter(**kwargs)
        
        if order_by:
            queryset = queryset.order_by(*order_by)
            
        if offset:
            queryset = queryset.offset(offset)
            
        if limit:
            queryset = queryset.limit(limit)
            
        return await queryset
        
    async def count(self, model_class: Type[Model], **kwargs) -> int:
        """
        统计记录数量
        
        Args:
            model_class: 模型类
            **kwargs: 过滤条件
            
        Returns:
            记录数量
        """
        self._ensure_initialized()
        return await model_class.filter(**kwargs).count()
        
    async def exists(self, model_class: Type[Model], **kwargs) -> bool:
        """
        检查记录是否存在
        
        Args:
            model_class: 模型类
            **kwargs: 查询条件
            
        Returns:
            是否存在
        """
        self._ensure_initialized()
        return await model_class.filter(**kwargs).exists()
        
    # 批量操作
    async def bulk_create(self, model_class: Type[Model], 
                         objects: List[Dict], batch_size: int = 1000) -> List[Model]:
        """
        批量创建记录
        
        Args:
            model_class: 模型类
            objects: 对象数据列表
            batch_size: 批次大小
            
        Returns:
            创建的模型实例列表
        """
        self._ensure_initialized()
        instances = [model_class(**obj) for obj in objects]
        created_instances = await model_class.bulk_create(instances, batch_size=batch_size)
        # Tortoise ORM 的 bulk_create 可能返回 None，我们返回原始实例
        return created_instances if created_instances is not None else instances
        
    async def bulk_update(self, instances: List[Model], 
                         fields: List[str], batch_size: int = 1000) -> None:
        """
        批量更新记录
        
        Args:
            instances: 模型实例列表
            fields: 要更新的字段列表
            batch_size: 批次大小
        """
        self._ensure_initialized()
        if instances:
            model_class = type(instances[0])
            await model_class.bulk_update(instances, fields, batch_size=batch_size)
            
    # 事务管理
    async def transaction(self, func: Callable, *args, **kwargs) -> Any:
        """
        在事务中执行函数
        
        Args:
            func: 要执行的异步函数
            *args: 函数参数
            **kwargs: 函数关键字参数
            
        Returns:
            函数返回值
        """
        self._ensure_initialized()
        async with in_transaction():
            return await func(*args, **kwargs)
            
    # 查询构建器
    def query(self, model_class: Type[Model]) -> 'QueryBuilder':
        """
        创建查询构建器
        
        Args:
            model_class: 模型类
            
        Returns:
            查询构建器实例
        """
        self._ensure_initialized()
        return QueryBuilder(model_class)
        
    # 原始 SQL 执行
    async def execute_raw(self, sql: str, params: List = None) -> List[Dict]:
        """
        执行原始 SQL 查询
        
        Args:
            sql: SQL 语句
            params: 参数列表
            
        Returns:
            查询结果
        """
        self._ensure_initialized()
        connection = connections.get("default")
        return await connection.execute_query_dict(sql, params or [])
        
    # 数据序列化
    def to_dict(self, instance: Model, exclude: List[str] = None, 
                include: List[str] = None) -> Dict:
        """
        将模型实例转换为字典
        
        Args:
            instance: 模型实例
            exclude: 排除的字段
            include: 包含的字段
            
        Returns:
            字典数据
        """
        data = {}
        fields = instance._meta.fields_map.keys()
        
        if include:
            fields = [f for f in fields if f in include]
        if exclude:
            fields = [f for f in fields if f not in exclude]
            
        for field in fields:
            value = getattr(instance, field, None)
            if hasattr(value, 'isoformat'):  # 处理日期时间
                value = value.isoformat()
            data[field] = value
            
        return data
        
    def to_dict_list(self, instances: List[Model], 
                     exclude: List[str] = None, 
                     include: List[str] = None) -> List[Dict]:
        """
        将模型实例列表转换为字典列表
        
        Args:
            instances: 模型实例列表
            exclude: 排除的字段
            include: 包含的字段
            
        Returns:
            字典数据列表
        """
        return [self.to_dict(instance, exclude, include) for instance in instances]
        
    # 实用工具方法
    async def truncate_table(self, model_class: Type[Model]) -> None:
        """
        清空表数据
        
        Args:
            model_class: 模型类
        """
        self._ensure_initialized()
        await model_class.all().delete()
        logger.debug(f"清空表: {model_class._meta.db_table}")
        
    async def drop_and_create_tables(self) -> None:
        """
        删除并重新创建所有表
        """
        self._ensure_initialized()
        await Tortoise.generate_schemas(safe=False)
        logger.info("重新创建所有表")
        
    async def get_table_info(self, model_class: Type[Model]) -> Dict:
        """
        获取表信息
        
        Args:
            model_class: 模型类
            
        Returns:
            表信息字典
        """
        self._ensure_initialized()
        meta = model_class._meta
        return {
            "table_name": meta.db_table,
            "fields": list(meta.fields_map.keys()),
            "pk_field": meta.pk.model_field_name if meta.pk else None,
            "db_table": meta.db_table,
            "fk_fields": list(meta.fk_fields),
            "m2m_fields": list(meta.m2m_fields)
        }


class QueryBuilder:
    """查询构建器类"""
    
    def __init__(self, model_class: Type[Model]):
        self.model_class = model_class
        self.queryset = model_class.all()
        
    def filter(self, **kwargs) -> 'QueryBuilder':
        """添加过滤条件"""
        self.queryset = self.queryset.filter(**kwargs)
        return self
        
    def exclude(self, **kwargs) -> 'QueryBuilder':
        """添加排除条件"""
        self.queryset = self.queryset.exclude(**kwargs)
        return self
        
    def order_by(self, *fields) -> 'QueryBuilder':
        """添加排序"""
        self.queryset = self.queryset.order_by(*fields)
        return self
        
    def limit(self, limit: int) -> 'QueryBuilder':
        """限制数量"""
        self.queryset = self.queryset.limit(limit)
        return self
        
    def offset(self, offset: int) -> 'QueryBuilder':
        """设置偏移量"""
        self.queryset = self.queryset.offset(offset)
        return self
        
    def select_related(self, *fields) -> 'QueryBuilder':
        """预加载关联对象"""
        self.queryset = self.queryset.select_related(*fields)
        return self
        
    def prefetch_related(self, *fields) -> 'QueryBuilder':
        """预取关联对象"""
        self.queryset = self.queryset.prefetch_related(*fields)
        return self
        
    async def all(self) -> List[Model]:
        """获取所有结果"""
        return await self.queryset
        
    async def first(self) -> Optional[Model]:
        """获取第一个结果"""
        return await self.queryset.first()
        
    async def get(self) -> Model:
        """获取单个结果"""
        return await self.queryset.get()
        
    async def count(self) -> int:
        """统计数量"""
        return await self.queryset.count()
        
    async def exists(self) -> bool:
        """检查是否存在"""
        return await self.queryset.exists()
        
    async def delete(self) -> int:
        """删除匹配的记录"""
        return await self.queryset.delete()
        
    async def update(self, **kwargs) -> int:
        """更新匹配的记录"""
        return await self.queryset.update(**kwargs)


# 全局实例
simple_orm = SimpleOrm()


# 装饰器
def with_db(db_url: str = None, models: List[str] = None, 
           generate_schemas: bool = False, auto_discover: bool = True):
    """
    数据库连接装饰器
    
    Args:
        db_url: 数据库连接 URL
        models: 模型模块列表，如果为 None 且 auto_discover=True 则自动发现
        generate_schemas: 是否自动生成数据库表结构
        auto_discover: 是否自动发现模型类
    """
    def decorator(func):
        async def wrapper(*args, **kwargs):
            if not simple_orm._initialized:
                await simple_orm.init(db_url, models, generate_schemas, auto_discover)
            try:
                return await func(*args, **kwargs)
            finally:
                pass  # 保持连接开启，由用户手动关闭
        return wrapper
    return decorator 